# test_prompt_improver.py
# Unit tests for prompt improver

import pytest
from unittest.mock import Mock, patch, MagicMock, AsyncMock
from typing import Dict, Any, List
import json
import asyncio

from tldw_Server_API.app.core.Prompt_Management.prompt_studio.prompt_improver import (
    PromptImprover, ImprovementStrategy, ImprovementResult, PromptAnalysis
)
from tldw_Server_API.app.core.DB_Management.PromptStudioDatabase import PromptStudioDatabase

########################################################################################################################
# Test PromptAnalysis Model

class TestPromptAnalysis:
    """Test the PromptAnalysis model."""

    def test_analysis_creation(self):
        """Test creating a PromptAnalysis."""
        analysis = PromptAnalysis(
            clarity_score=0.85,
            specificity_score=0.75,
            structure_score=0.90,
            completeness_score=0.80,
            issues=["Too verbose", "Missing examples"],
            suggestions=["Simplify language", "Add few-shot examples"],
            metrics={"avg_sentence_length": 25, "complexity": 0.7}
        )

        assert analysis.clarity_score == 0.85
        assert len(analysis.issues) == 2
        assert "complexity" in analysis.metrics

    def test_analysis_overall_score(self):
        """Test calculating overall score."""
        analysis = PromptAnalysis(
            clarity_score=0.80,
            specificity_score=0.70,
            structure_score=0.90,
            completeness_score=0.60
        )

        overall = analysis.overall_score()
        assert overall == 0.75  # Average of all scores

    def test_analysis_to_dict(self):
        """Test converting analysis to dictionary."""
        analysis = PromptAnalysis(
            clarity_score=0.85,
            specificity_score=0.75,
            issues=["Issue 1"],
            suggestions=["Suggestion 1"]
        )

        result = analysis.to_dict()
        assert isinstance(result, dict)
        assert result["clarity_score"] == 0.85
        assert len(result["issues"]) == 1

########################################################################################################################
# Test ImprovementResult Model

class TestImprovementResult:
    """Test the ImprovementResult model."""

    def test_result_creation(self):
        """Test creating an ImprovementResult."""
        result = ImprovementResult(
            original_prompt="Original text",
            improved_prompt="Improved text",
            strategy=ImprovementStrategy.CLARITY,
            improvements_made=["Simplified language", "Fixed grammar"],
            score_change=0.15,
            metadata={"iterations": 3}
        )

        assert result.original_prompt == "Original text"
        assert result.strategy == ImprovementStrategy.CLARITY
        assert result.score_change == 0.15
        assert result.metadata["iterations"] == 3

    def test_result_with_analysis(self):
        """Test result with before/after analysis."""
        before = PromptAnalysis(
            clarity_score=0.60,
            specificity_score=0.60,
            structure_score=0.60,
            completeness_score=0.60
        )
        after = PromptAnalysis(
            clarity_score=0.85,
            specificity_score=0.85,
            structure_score=0.85,
            completeness_score=0.85
        )

        result = ImprovementResult(
            original_prompt="Original",
            improved_prompt="Improved",
            strategy=ImprovementStrategy.CLARITY,
            before_analysis=before,
            after_analysis=after
        )

        improvement = result.calculate_improvement()
        assert improvement == 0.25

########################################################################################################################
# Test PromptImprover

class TestPromptImprover:
    """Test the PromptImprover class."""

    @pytest.fixture
    def improver(self):
        """Create a PromptImprover instance."""
        mock_db = Mock(spec=PromptStudioDatabase)
        mock_db.client_id = "test-client"
        return PromptImprover(db=mock_db)

    def test_improver_initialization(self):
        """Test PromptImprover initialization."""
        mock_db = Mock(spec=PromptStudioDatabase)
        mock_db.client_id = "test-client"
        improver = PromptImprover(db=mock_db)
        assert hasattr(improver, 'db')
        assert improver.db == mock_db

    def test_analyze_prompt(self, improver):
        """Test analyzing a prompt."""
        prompt = "Please summarize the following text in a concise manner"

        analysis = improver.analyze(prompt)

        assert isinstance(analysis, PromptAnalysis)
        assert 0 <= analysis.clarity_score <= 1
        assert 0 <= analysis.specificity_score <= 1
        assert isinstance(analysis.issues, list)
        assert isinstance(analysis.suggestions, list)

    def test_improve_clarity(self, improver):
        """Test improving prompt clarity."""
        original = "You should maybe try to perhaps summarize this text if you can"

        result = improver.improve(
            prompt=original,
            strategy=ImprovementStrategy.CLARITY
        )

        assert result.improved_prompt != original
        assert len(result.improvements_made) > 0
        assert result.strategy == ImprovementStrategy.CLARITY

    def test_improve_specificity(self, improver):
        """Test improving prompt specificity."""
        original = "Do something with this data"

        result = improver.improve(
            prompt=original,
            strategy=ImprovementStrategy.SPECIFICITY,
            context={"task": "analysis", "data_type": "sales"}
        )

        assert "data" in result.improved_prompt
        # Should be more specific than original
        assert len(result.improved_prompt) > len(original)

    def test_improve_structure(self, improver):
        """Test improving prompt structure."""
        original = "Analyze this. Be thorough. Use examples. Format nicely."

        result = improver.improve(
            prompt=original,
            strategy=ImprovementStrategy.STRUCTURE
        )

        # Should have better structure
        assert result.improved_prompt != original
        assert any(word in result.improved_prompt.lower()
                  for word in ["first", "then", "finally", "step"])

    def test_add_examples(self, improver):
        """Test adding examples to prompt."""
        original = "Classify the sentiment of the text"
        examples = [
            {"input": "I love this!", "output": "positive"},
            {"input": "This is terrible", "output": "negative"}
        ]

        result = improver.improve(
            prompt=original,
            strategy=ImprovementStrategy.ADD_EXAMPLES,
            examples=examples
        )

        assert "I love this!" in result.improved_prompt
        assert "positive" in result.improved_prompt

    def test_add_constraints(self, improver):
        """Test adding constraints to prompt."""
        original = "Generate a summary"
        constraints = [
            "Maximum 100 words",
            "Include key points only",
            "Use bullet points"
        ]

        result = improver.improve(
            prompt=original,
            strategy=ImprovementStrategy.ADD_CONSTRAINTS,
            constraints=constraints
        )

        assert "100 words" in result.improved_prompt
        assert "bullet" in result.improved_prompt.lower()

    def test_chain_of_thought(self, improver):
        """Test adding chain-of-thought reasoning."""
        original = "Solve this math problem: 25 * 4 + 10"

        result = improver.improve(
            prompt=original,
            strategy=ImprovementStrategy.CHAIN_OF_THOUGHT
        )

        cot_keywords = ["step by step", "think through", "reason", "explain"]
        assert any(keyword in result.improved_prompt.lower() for keyword in cot_keywords)

    def test_improve_multiple_strategies(self, improver):
        """Test applying multiple improvement strategies."""
        original = "Analyze data"

        result = improver.improve_multi(
            prompt=original,
            strategies=[
                ImprovementStrategy.CLARITY,
                ImprovementStrategy.SPECIFICITY,
                ImprovementStrategy.STRUCTURE
            ]
        )

        # Should be significantly improved
        assert len(result.improved_prompt) > len(original) * 2
        assert len(result.improvements_made) >= 3

    def test_auto_improve(self, improver):
        """Test automatic improvement selection."""
        original = "do stuff with things"

        result = improver.auto_improve(original)

        # Auto improve should detect issues and apply appropriate strategies
        assert result.improved_prompt != original
        assert len(result.improvements_made) > 0
        # The strategy may be changed to the selected one (e.g., SPECIFICITY)
        assert result.strategy in [ImprovementStrategy.AUTO, ImprovementStrategy.SPECIFICITY]

    def test_iterative_improvement(self, improver):
        """Test iterative prompt improvement."""
        original = "Summarize"

        results = improver.improve_iterative(
            prompt=original,
            max_iterations=3,
            target_score=0.8
        )

        assert len(results) <= 3
        # Each iteration should be different
        prompts = [r.improved_prompt for r in results]
        assert len(set(prompts)) == len(prompts)

    @pytest.mark.asyncio
    async def test_async_improvement(self, improver):
        """Test asynchronous prompt improvement."""
        original = "Analyze this text"

        result = await improver.improve_async(
            prompt=original,
            strategy=ImprovementStrategy.CLARITY
        )

        assert result.improved_prompt != original

    def test_batch_improvement(self, improver):
        """Test batch prompt improvement."""
        prompts = [
            "Do analysis",
            "Make summary",
            "Write code"
        ]

        results = improver.improve_batch(
            prompts=prompts,
            strategy=ImprovementStrategy.CLARITY
        )

        assert len(results) == 3
        assert all(r.improved_prompt != prompts[i] for i, r in enumerate(results))

    def test_compare_prompts(self, improver):
        """Test comparing two prompts."""
        prompt1 = "Summarize this text"
        prompt2 = "Please provide a comprehensive summary of the following text, including key points and conclusions"

        comparison = improver.compare(prompt1, prompt2)

        assert "prompt1" in comparison
        assert "prompt2" in comparison
        assert "winner" in comparison
        assert comparison["winner"] == "prompt2"  # More detailed prompt should win

    def test_validate_improvement(self, improver):
        """Test validating prompt improvement."""
        original = "Do task"
        improved = "Please complete the following task with attention to detail"

        is_valid = improver.validate_improvement(original, improved)

        assert is_valid is True

    def test_improvement_with_llm(self, improver):
        """Test improvement using LLM assistance."""
        with patch.object(improver, 'llm_client') as mock_llm:
            mock_llm.generate.return_value = "Improved prompt from LLM"

            original = "Basic prompt"
            result = improver.improve_with_llm(
                prompt=original,
                model="gpt-4"
            )

            assert result.improved_prompt == "Improved prompt from LLM"
            assert mock_llm.generate.called

    def test_improve_text_parses_openai_response(self, improver, monkeypatch):
        """Ensure _improve_text handles OpenAI-format responses."""
        payload = {
            "id": "chatcmpl-test",
            "object": "chat.completion",
            "created": 123,
            "model": "gpt-4o-mini",
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": "Improved prompt content"
                    }
                }
            ]
        }

        monkeypatch.setattr(
            "tldw_Server_API.app.core.Prompt_Management.prompt_studio.prompt_improver.chat_with_openai",
            lambda **kwargs: payload,
        )

        improved = improver._improve_text(
            text="Original prompt",
            strategies=["clarity"],
            text_type="prompt",
            model_name="gpt-4o-mini",
        )

        assert improved == "Improved prompt content"

########################################################################################################################
# Test Improvement Strategies

class TestImprovementStrategies:
    """Test individual improvement strategies."""

    @pytest.fixture
    def improver(self):
        """Create improver instance."""
        mock_db = Mock(spec=PromptStudioDatabase)
        mock_db.client_id = "test-client"
        return PromptImprover(db=mock_db)

    def test_simplify_strategy(self, improver):
        """Test simplification strategy."""
        complex_prompt = """
        In order to facilitate the comprehensive analysis of the provided textual data,
        please endeavor to utilize your computational capabilities to generate a
        condensed representation of the salient information contained therein.
        """

        result = improver.improve(
            prompt=complex_prompt,
            strategy=ImprovementStrategy.SIMPLIFY
        )

        # Should be shorter and simpler
        assert len(result.improved_prompt) < len(complex_prompt)
        # Check for simpler language
        assert "facilitate" not in result.improved_prompt
        assert "endeavor" not in result.improved_prompt

    def test_expand_strategy(self, improver):
        """Test expansion strategy."""
        brief_prompt = "Summarize"

        result = improver.improve(
            prompt=brief_prompt,
            strategy=ImprovementStrategy.EXPAND
        )

        # Should be longer and more detailed
        assert len(result.improved_prompt) > len(brief_prompt) * 5

    def test_formalize_strategy(self, improver):
        """Test formalization strategy."""
        informal_prompt = "hey can u plz check this code and tell me if its ok?"

        result = improver.improve(
            prompt=informal_prompt,
            strategy=ImprovementStrategy.FORMALIZE
        )

        # Should be more formal - check that at least one informal term was replaced
        improved_lower = result.improved_prompt.lower()
        formal_improvements = (
            "hey" not in improved_lower or
            " u " not in result.improved_prompt or
            "plz" not in result.improved_prompt
        )
        assert formal_improvements

    def test_technical_strategy(self, improver):
        """Test technical enhancement strategy."""
        general_prompt = "Check if the function works correctly"

        result = improver.improve(
            prompt=general_prompt,
            strategy=ImprovementStrategy.TECHNICAL,
            context={"language": "python", "task": "unit testing"}
        )

        # Should include technical terms
        technical_terms = ["test", "assert", "edge case", "exception", "validation"]
        assert any(term in result.improved_prompt.lower() for term in technical_terms)

########################################################################################################################
# Test Analysis Features

class TestAnalysisFeatures:
    """Test prompt analysis features."""

    @pytest.fixture
    def improver(self):
        """Create improver instance."""
        mock_db = Mock(spec=PromptStudioDatabase)
        mock_db.client_id = "test-client"
        return PromptImprover(db=mock_db)

    def test_readability_analysis(self, improver):
        """Test readability analysis."""
        simple = "Write a summary of this text."
        complex = "Endeavor to synthesize a comprehensive yet succinct representation."

        simple_score = improver.analyze_readability(simple)
        complex_score = improver.analyze_readability(complex)

        assert simple_score > complex_score

    def test_ambiguity_detection(self, improver):
        """Test detecting ambiguous language."""
        ambiguous = "Do the thing with the stuff over there"
        clear = "Process the data file located in the input directory"

        ambiguous_issues = improver.detect_ambiguity(ambiguous)
        clear_issues = improver.detect_ambiguity(clear)

        assert len(ambiguous_issues) > len(clear_issues)

    def test_completeness_check(self, improver):
        """Test checking prompt completeness."""
        incomplete = "Analyze"
        complete = "Analyze the provided sales data for Q4 2023, focusing on trends, anomalies, and year-over-year growth"

        incomplete_score = improver.check_completeness(incomplete)
        complete_score = improver.check_completeness(complete)

        assert complete_score > incomplete_score

    def test_task_alignment(self, improver):
        """Test checking task alignment."""
        prompt = "Summarize the key findings from the research paper"
        task = "summarization"

        alignment_score = improver.check_task_alignment(prompt, task)

        # Should have some alignment since "Summarize" is in prompt
        # but the exact score depends on implementation
        assert alignment_score >= 0.0  # Any alignment

########################################################################################################################
# Test Error Handling

class TestErrorHandling:
    """Test error handling in PromptImprover."""

    @pytest.fixture
    def improver(self):
        """Create improver instance."""
        mock_db = Mock(spec=PromptStudioDatabase)
        mock_db.client_id = "test-client"
        return PromptImprover(db=mock_db)

    def test_empty_prompt_handling(self, improver):
        """Test handling empty prompt."""
        with pytest.raises(ValueError):
            improver.improve("", strategy=ImprovementStrategy.CLARITY)

    def test_invalid_strategy(self, improver):
        """Test handling invalid strategy."""
        # Invalid strategy should either raise error or return unchanged
        try:
            result = improver.improve("Test prompt", strategy="invalid_strategy")
            # If no error, check that prompt is unchanged
            assert result.improved_prompt == "Test prompt"
        except (ValueError, AttributeError, TypeError):
            # Expected error
            pass

    def test_malformed_examples(self, improver):
        """Test handling malformed examples."""
        # Test with malformed examples - should handle gracefully
        try:
            result = improver.improve(
                "Test prompt",
                strategy=ImprovementStrategy.ADD_EXAMPLES,
                examples="not a list"
            )
            # If no error, check result
            assert result.improved_prompt == "Test prompt"
        except (ValueError, AttributeError, TypeError):
            # Expected error
            pass

    def test_improvement_failure_recovery(self, improver):
        """Test recovery from improvement failure."""
        # Simulate a strategy that fails
        with patch.object(improver, '_apply_strategy') as mock_apply:
            mock_apply.side_effect = Exception("Strategy failed")

            original = "Test prompt"
            result = improver.improve(
                prompt=original,
                strategy=ImprovementStrategy.CLARITY,
                fallback=True
            )

            # Should return original prompt on failure with fallback
            assert result.improved_prompt == original
            assert "error" in result.metadata

########################################################################################################################
# Test Caching and Performance

class TestCachingPerformance:
    """Test caching and performance features."""

    @pytest.fixture
    def improver(self):
        """Create improver with caching enabled."""
        mock_db = Mock(spec=PromptStudioDatabase)
        mock_db.client_id = "test-client"
        return PromptImprover(db=mock_db)

    def test_result_caching(self, improver):
        """Test caching of improvement results."""
        prompt = "Analyze this data"

        # First call
        result1 = improver.improve(prompt, strategy=ImprovementStrategy.CLARITY)

        # Second call with same parameters
        result2 = improver.improve(prompt, strategy=ImprovementStrategy.CLARITY)

        # Should return cached result
        assert result1.improved_prompt == result2.improved_prompt

    def test_cache_invalidation(self, improver):
        """Test cache invalidation."""
        prompt = "Test prompt"

        # Cache a result
        result1 = improver.improve(prompt, strategy=ImprovementStrategy.CLARITY)

        # Invalidate cache
        improver.clear_cache()

        # Should generate new result
        with patch.object(improver, '_apply_strategy') as mock_apply:
            mock_apply.return_value = "Different improved prompt"
            result2 = improver.improve(prompt, strategy=ImprovementStrategy.CLARITY)

            assert result2.improved_prompt != result1.improved_prompt

    def test_batch_processing_performance(self, improver):
        """Test batch processing is more efficient than individual."""
        import time

        prompts = ["Prompt " + str(i) for i in range(10)]

        # Individual processing
        start = time.time()
        individual_results = [
            improver.improve(p, strategy=ImprovementStrategy.CLARITY)
            for p in prompts
        ]
        individual_time = time.time() - start

        # Batch processing
        start = time.time()
        batch_results = improver.improve_batch(
            prompts,
            strategy=ImprovementStrategy.CLARITY
        )
        batch_time = time.time() - start

        # Batch should be faster (or at least not significantly slower)
        assert batch_time <= individual_time * 1.1

########################################################################################################################
# Test Integration Features

class TestIntegrationFeatures:
    """Test integration with other components."""

    @pytest.fixture
    def improver(self):
        """Create improver instance."""
        mock_db = Mock(spec=PromptStudioDatabase)
        mock_db.client_id = "test-client"
        return PromptImprover(db=mock_db)

    def test_export_improvements(self, improver):
        """Test exporting improvement history."""
        prompts = ["Prompt 1", "Prompt 2", "Prompt 3"]

        results = []
        for prompt in prompts:
            result = improver.improve(prompt, strategy=ImprovementStrategy.CLARITY)
            results.append(result)

        # Export to JSON
        exported = improver.export_improvements(results, format="json")
        data = json.loads(exported)

        assert len(data) == 3
        assert all("original_prompt" in item for item in data)
        assert all("improved_prompt" in item for item in data)

    def test_import_strategies(self, improver):
        """Test importing custom strategies."""
        custom_strategies = {
            "custom_1": lambda p: p.upper(),
            "custom_2": lambda p: p + " [ENHANCED]"
        }

        improver.import_strategies(custom_strategies)

        # The current implementation may not directly use custom strategies
        # So we just test that import doesn't fail
        assert hasattr(improver, 'custom_strategies')
        assert improver.custom_strategies == custom_strategies

    def test_improvement_pipeline(self, improver):
        """Test creating improvement pipeline."""
        pipeline = improver.create_pipeline([
            ImprovementStrategy.CLARITY,
            ImprovementStrategy.STRUCTURE,
            ImprovementStrategy.ADD_CONSTRAINTS
        ])

        original = "Do analysis"
        result = pipeline.run(original, constraints=["Be concise", "Use examples"])

        # Should apply all strategies
        assert len(result.improvements_made) >= 3
        assert result.improved_prompt != original
